#version 450 compatibility





#define BlockNum 16

layout(local_size_x = DEF_LocalSize, local_size_y = DEF_LocalSize, local_size_z = DEF_LocalSize) in;

layout(binding = 0) uniform Buf_GlobalParameter0 {
    uint m_ElementNum;      //元素数量
    uint m_BitStep;         //比特位偏移
    uint m_MergeOffset;     //合并偏移
    uint m_Section;         //排序阶段
    uint m_ThreadNum;       //线程数
};

layout(std430, binding = 1) buffer Buf_GlobalParameter1 {
    uvec2 m_Hash[];
};
layout(std430, binding = 2) buffer Buf_GlobalParameter2 {
    uvec2 m_BarrelElementOffset[];
};
layout(std430, binding = 3) buffer Buf_GlobalParameter3 {
    uint m_BarrelElementNum[];
};



//----------------桶分配------------------------
void f_BarrelDistribution(in uint id);

//----------------元素偏移----------------------
void f_ElementOffset(in uint id);
void f_ElementOffsetCopy(in uint id);

//----------------重新排序----------------------
void f_Rearrange(in uint id);



void main(){
    const uint id = (gl_WorkGroupID.y * gl_NumWorkGroups.x) + gl_WorkGroupID.x;

    if(id<m_ThreadNum) {
        switch(m_Section){
        case 0:
            f_BarrelDistribution(id);
            break;
        case 1:
            f_ElementOffset(id);
            break;
        case 2:
            f_ElementOffsetCopy(id);
            break;
        case 3:
            f_Rearrange(id);
            break;
        }
    }
    

}






void f_BarrelDistribution(in uint id){
    const uint HashOffset = id * BlockNum;

    uint globalID = 0;
    uint tmpHash[16]; 
    uint newHash[16]; 
    uint countTmp[16];

    bool mask[16] = bool[16](false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false);

    uint ic = 0;
    for (uint i=0; i<16; ++i) {
        if(HashOffset+i < m_ElementNum){
            tmpHash[i] = m_Hash[HashOffset+i].x;
            ++ic;
        } else {
            tmpHash[i] = 0xffffffff;
        }
        
        countTmp[i] = 0;
    }

    for (uint i = 0; i<16; ++i) {
        for (uint ia = 0; ia < ic; ++ia) {
            if (mask[ia] || tmpHash[ia].x == 0xffffffff) {
                continue;
            }
            uint hash = bitfieldExtract(tmpHash[ia].x, int(m_BitStep * 4), 4);

            if (i == hash) {
                newHash[globalID] = tmpHash[ia];

                ++countTmp[i];
                ++globalID;
                mask[ia] = true;

                if(globalID >= ic) break;
            }
        }
        //if(globalID >= ic) break;
    }

    for(uint i=0; i<16; ++i){
        m_BarrelElementOffset[HashOffset+i] = uvec2(countTmp[i]);
        m_BarrelElementNum[HashOffset+i] = countTmp[i];
        
        m_Hash[HashOffset + i].y = newHash[i];
    }
    //m_BarrelElementNum[0] = 0;
}

void f_ElementOffset(in uint id){
    const uint HashOffset = id * BlockNum;
    uint offset = (id + m_MergeOffset);
    if(offset < m_ThreadNum){
        offset *= BlockNum;
        for(int i=0; i<BlockNum; ++i){
            m_BarrelElementOffset[offset].y += m_BarrelElementOffset[HashOffset + i].x;
            ++offset;
        }
    }
    
}
void f_ElementOffsetCopy(in uint id){
    uint offset = (id + m_MergeOffset);
    if(offset < m_ThreadNum){
        offset *= BlockNum;
        for(int i=0; i<BlockNum; ++i){
            m_BarrelElementOffset[offset].x = m_BarrelElementOffset[offset].y;
            ++offset;
        }
    }
    
}


void f_Rearrange(in uint id){
    //每个进制取的偏移 0~1 1~2 。。。。 14~15
    uint totalOffset = 0;

    uint currentOffset = id*BlockNum;
    uint currentOffset_s1 = id*BlockNum;
    uint endPos = (m_ThreadNum-1)*BlockNum;
    
    for(uint i=0; i<16; ++i){
        uint RangeSize = m_BarrelElementNum[currentOffset];
        
        if(RangeSize>0){
            uint currentStep = m_BarrelElementOffset[currentOffset].x - RangeSize + totalOffset;
            for (uint ia = 0; ia < RangeSize; ++ia) {
                m_Hash[currentStep + ia].x = m_Hash[currentOffset_s1 + ia].y;
            }
            currentOffset_s1 += RangeSize;
        }
        
        totalOffset += m_BarrelElementOffset[endPos+i].x;
        ++currentOffset;
    }


}



